#!/usr/bin/env python

#Original Script by RiCON; 
#http://forum.doom9.org/showthread.php?t=154535

#base on vfr.py 0.6
#modify by cunhan
#v0.2
#add: Adjust assfile extracted from ts according to trims
#v0.1
#add: in order to output aac file, using besplit to split audio

import sys
import re
import linecache
import optparse
import os
import random
import datetime
import win32api
from subprocess import call

rat = re.compile('(\d+)(?:/|:)(\d+)')
v1re = re.compile('# timecode format v1')
v2re = re.compile('# timecode format v2')
trimre = re.compile("(?<!#)Trim\((\d+)\s*,\s*(\d+)\)",re.I)
tcConv = 'tcConv'
mkvmerge = 'mkvmerge'
beSplit = 'beSplit'

def main():
    
    linecache.clearcache()
    
    p = optparse.OptionParser(description='Grabs avisynth trims and outputs chapter file, qpfile and/or cuts audio (works with cfr and vfr input)',
                              prog='TrimApp.py',
                              version='VFR Chapter Creator Mod 0.1',
                              usage='%prog [options] infile.avs')
    p.add_option('--input', '-i', action="store", help='Audio file to be cut', dest="input")
    p.add_option('--output', '-o', action="store", help='Cut audio from MKVMerge', dest="output")
    p.add_option('--fps', '-f', action="store", help='Frames per second (for cfr input)', dest="fps")
    p.add_option('--timecodes', '-t', action="store", help='Timecodes file from the vfr video (v1 needs tcConv)', dest="timecodes")
    p.add_option('--chapters', '-c', action="store", help='Chapters file', dest="chapters")
    p.add_option('--qpfile', '-q', action="store", help='QPFile for x264', dest="qpfile")
    p.add_option('--verbose', '-v', action="store_true", help='Verbose', dest="verbose")
    p.add_option('--merge', '-m', action="store_true", help='Merge cut files', dest="merge")
    p.add_option('--remove', '-r', action="store_true", help='Remove cut files', dest="remove")
    p.add_option('--frames', action="store", help='Number of frames for v1 conversion', dest="frames")
    p.add_option('--test', action="store_true", help="Test mode (do not create new files)", dest="test")
    p.add_option('--beSplit', '-b', action="store_true", help='using BeSplit to Cut Audio file', dest="beSplit")
    p.add_option('--ass', action="store", help='Ajust ASS timestamp by trims', dest="ass")
    (options, args) = p.parse_args()
    
    if len(args) != 1:
        p.error("No avisynth script specified.")
    elif options.timecodes == None and os.path.isfile(args[0] + ".tc.txt") == False and options.fps == None:
        options.timecodes = '30000/1001'
    elif options.timecodes == None and os.path.isfile(args[0] + ".tc.txt") == True:
        options.timecodes = args[0] + ".tc.txt"
    elif options.timecodes != None and options.fps != None:
        p.error("Can't use vfr input AND cfr input")
    elif options.timecodes != None and os.path.isfile(options.timecodes) == True:
        options.timecodes = options.timecodes
    else:
        options.timecodes = options.fps
    
    if options.chapters != None:
        chapType = 'MKV' if options.chapters[-3:] == 'xml' else 'OGM'
    else:
        chapType = ''
    
    if options.output == None and options.input != None:
        options.output = '%s.%s.mka' % (options.input, args[0][:-4])
        #get output file extension, to see it's aac or ac3 or mka
        bExt = options.output[-3:]
    
    
    
        
    if options.verbose == True:
        status = """
Avisynth file:   {input}
Audio file:      {audio}
Cut Audio file:  {cutaudio}
Timecodes/FPS:   {timecodes}
Chapters file:   {chapters}
QP file:         {qpfile}

Merge/Rem files: {merge}/{remove}
Verbose:         {verbose}
Test Mode:       {test}
""".format(input=args[0],
            audio=options.input,
            cutaudio=options.output,
            timecodes=options.timecodes,
            chapters=options.chapters,
            qpfile=options.qpfile,
            merge=options.merge,
            remove=options.remove,
            verbose=options.verbose,
            test=options.test)
        #print(status)
    
    quiet = '' if options.verbose == True else '-q'
    audio = []
    Trims = []
    
    with open(args[0], "r") as avs:
        # use only the first non-commented line with trims
        for line in avs:
            if trimre.match(line) != None:
                Trims = trimre.findall(line)
                break
            
        if len(Trims) < 1:
            sys.exit("Error: Avisynth script has no uncommented trims")
        
        if options.verbose == True:
            print('In trims:  %s'%(', '.join(['(%s,%s)'%(i[0],i[1]) for i in Trims])))
        
        # trims' offset calculation
        Trims2 = []
        Trims2ts = []
        options.timecodes = [options.timecodes, determineFormat(options.timecodes)]
        if options.timecodes[1] == 2:
            tc = options.timecodes
            nTrims = int(options.frames) if options.frames != None else int(Trims[-1][1])+2
            if os.path.isfile(tc[0]+"v2.txt") == False:
                tcConv = call('"%s" "%s" "%s" %d' % (tcConv, tc[0], tc[0]+"v2.txt", nTrims))
                if tcConv > 0:
                    sys.exit("Failed to execute tcConv: %d; Please put it in your path" % tcConv)
            options.timecodes[0] = tc[0]+"v2.txt"
        for i in range(len(Trims)):
        
            fn1 = int(Trims[i][0])                      # first frame
            fn1ts = Ts(fn1,options.timecodes)         # first frame timecode
            fn2 = int(Trims[i][1])                      # last frame
            fn2ts = Ts(fn2,options.timecodes)         # last frame timecode
            fn2tsaud = Ts(fn2+1,options.timecodes)    # last frame timecode for audio
                        
            if i != 0:                                  # if it's not the first trim
                last = int(Trims[i-1][1])+1
                lastts = Ts(last,options.timecodes)
                offset += fn1-last
                offsetts += fn1ts-lastts
            elif fn1 > 0:                               # if the first trim doesn't start at 0
                offset = fn1
                offsetts = fn1ts
            else:
                offset = 0
                offsetts = 0
            
            # apply the offset to the trims
            Trims2.append([fn1-offset,fn2-offset])
            Trims2ts.append([fn1ts-offsetts,fn2ts-offsetts])
            
            # make list with timecodes to cut audio
            audio.append(fn1ts)
            audio.append(fn2tsaud)
    
    if options.verbose == True:
        print('Out trims: %s'%(', '.join(['(%d,%d)'%(i[0],i[1]) for i in Trims2])))
    
    # make qpfile
    if options.qpfile != None and options.test == None:
        with open(options.qpfile, "w") as qpf:
            for trim in Trims2:
                qpf.write('%s I -1\n' % trim[0])
    
    # adjust assfile extracted from ts according to trims
    if options.ass != None :
        #print audio
        if os.path.isfile(options.ass) == True :
            newasslines = []
            assfile = open(options.ass,"r")
            asslines = assfile.readlines()
            assfile.close()
            for line in asslines:
                if line.upper().startswith("DIALOGUE") == True :
                    tokens = line.split(",")
                    tokens[1] = adjustTimeStamp(tokens[1],audio)
                    tokens[2] = adjustTimeStamp(tokens[2],audio)
                    newline = ",".join(tokens)
                    newasslines.append(newline)
                else:
                    newasslines.append(line)
            newass = open(options.ass[:-4]+".adjust.ass","w")
            for line in newasslines:
                newass.write(line)
            newass.close()              
            
    # make audio cuts
    if options.input != None:
        print audio
        delay = re.search('DELAY ([-]?\d+)',options.input).group(1) if re.search('DELAY ([-]?\d+)',options.input) != None else '0'
        
        #using besplit
        if options.beSplit != None:
            cuttimes = ' '.join([str(a*1.0/1000000000) for a in audio])
            print audio
            print cuttimes
            cutCmd = '"%s" -core( -input "%s" -prefix "%s" -type %s -a ) -split( %s )'%\
            (beSplit, options.input, options.output + '.split', bExt, cuttimes )
            if options.verbose == True:
                print 'Cutting: %s\n' % cutCmd
            if options.test == None:
                cutExec = call(cutCmd)
                #when beSplit Complete,it retrun 58 but no 0. that's a problem
                if cutExec > 0 and cutExec != 58:
                    sys.exit("Failed to execute beSplit: %d" % cutExec)
            if options.merge == True:
                merge = []
                for i in range(1,len(audio)):
                    if (i % 2 != 0):
                        splitedFile = '%s.split%02d.%s' % (options.output, i, bExt)
                        splitedFile = win32api.GetShortPathName(splitedFile)
                        merge.append(splitedFile)
                listfile = os.path.join(os.path.split(options.input)[0],"join.txt")
                fh = open(listfile, "w")
                for l in merge:
                    fh.write(l + "\n")
                fh.close()
                mergeCmd = '"%s" -core( -input "%s" -prefix "%s" -join )' % \
                (beSplit,listfile, options.output)
                if options.verbose == True:
                    print 'Merging: %s\n' % ' '.join(merge)
                    print mergeCmd
                if options.test == None:
                    mergeExec = call(mergeCmd)
                    #
                    if mergeExec > 0 and mergeExec != 58:                        
                        sys.exit("Failed to execute besplit: %d" % mergeExec)
            if options.remove == True:
                remove = ['%s.split%02d.%s' % (options.output, i, bExt) for i in range(1,len(audio))]
                if options.verbose == True:
                    print('Removing: %s\n' % ', '.join(remove))
                if options.test == None:
                    try:
                        [os.unlink(i) for i in remove]
                    finally:
                        pass
                
        else:
            if audio[0] == 0:
                includefirst = True
                audio = audio[1:]
            else:
                includefirst = False
                
            audio = [formatTime(i) for i in audio]
            cuttimes = ','.join(audio)
            cutCmd = '"%s" -o "%s" "%s" --sync 0:%s --split timecodes:%s %s' % (mkvmerge, options.output + '.split.mka', options.input, delay, cuttimes, quiet)
            if options.verbose == True:
                print('Cutting: %s\n' % cutCmd)
            if options.test == None:
                cutExec = call(cutCmd)
                #it's strang that mergeExec return 1 but not 0, even if the process successed.
#                if cutExec > 0:
#                    sys.exit("Failed to execute mkvmerge: %d" % cutExec)
            if options.merge == True:
                merge = []
                for i in range(1,len(audio)+2):
                    if (includefirst == True and i % 2 != 0) or (includefirst == False and i % 2 == 0):
                        merge.append('+%s.split-%03d.mka' % (options.output, i))
                merge[0] = merge[0][1:]
                mergeCmd = '"%s" -o "%s" %s %s' % (mkvmerge,options.output, ' '.join(merge), quiet)
                if options.verbose == True:
                    print('Merging: %s\n' % ' '.join(merge))
                if options.test == None:
                    mergeExec = call(mergeCmd)
                    if mergeExec > 0:
                        sys.exit("Failed to execute mkvmerge: %d" % mergeExec)
            
            if options.remove == True:
                remove = ['%s.split-%03d.mka' % (options.output, i) for i in range(1,len(audio)+2)]
                if options.verbose == True:
                    print('Removing: %s\n' % ', '.join(remove))
                if options.test == None:
                    [os.unlink(i) for i in remove]
            
    EditionUID = random.randint(100000,1000000)
    matroskaXmlHeader = '<?xml version="1.0" encoding="UTF-8"?>\n<!-- <!DOCTYPE Tags SYSTEM "matroskatags.dtd"> -->\n<Chapters>'
    matroskaXmlEditionHeader = """
	<EditionEntry>
		<EditionFlagHidden>{0}</EditionFlagHidden>
		<EditionFlagDefault>{1}</EditionFlagDefault>
		<EditionFlagOrdered>{2}</EditionFlagOrdered>
		<EditionUID>{3}</EditionUID>
""".format(0,1,1,EditionUID)
    matroskaXmlEditionFooter = '	</EditionEntry>'
    matroskaXmlFooter = '\n</Chapters>'
    
    matroskaXmlTagsHeader = '<?xml version="1.0" encoding="UTF-8"?>\n<!-- <!DOCTYPE Tags SYSTEM "matroskatags.dtd"> -->\n<Tags>'
    matroskaXmlTagsEdition = """
	<Tag>
		<Targets>
			<EditionUID>{0}</EditionUID>
			<TargetTypeValue>50</TargetTypeValue>
		</Targets>

		<Simple>
			<Name>TITLE</Name>
			<String>{1}</String>
			<TagLanguage>{2}</TagLanguage>
			<DefaultLanguage>1</DefaultLanguage>
		</Simple>

	</Tag>""".format(EditionUID,"Default","eng")
    
    if chapType == 'OGM':
        matroskaXmlHeader = matroskaXmlEditionHeader = matroskaXmlEditionFooter = matroskaXmlFooter = ''
    
    if options.test == None and chapType != '':
        with open(options.chapters, "w") as output:
        
            output.write(matroskaXmlHeader)
            
            output.write(matroskaXmlEditionHeader)
            
            [output.write(generateChap(formatTime(Trims2ts[i][0]), formatTime(Trims2ts[i][1]),i+1,chapType)) for i in range(len(Trims2ts))]
            
            output.write(matroskaXmlEditionFooter)
            
            output.write(matroskaXmlFooter)
    elif chapType != '':
        print("Writing %s Chapters to %s"%(chapType,options.chapters))

def adjustTimeStamp(timestr, timeInterval):
    
    temp = []
    temp.extend(timeInterval)
    temp.append(0)
    t =datetime.datetime.strptime(timestr,"%H:%M:%S.%f")
    #print timestr, t
    timestamp = ((t.hour*3600 +t.minute*60 + t.second)*1000000 + t.microsecond)*1000
    #print timestr, t, timestamp
    offset = 0
    for i in range(len(temp)):
        if i%2 == 0:
            offset = offset + (temp[i]-temp[i-1])
        if timestamp <= temp[i]:
            adjustTime =  timestamp - offset
            strT = formatTime(adjustTime)
            strT = strT[0:strT.index('.')+3]
            return strT
    
    return "00:00:00.000"

def formatTime(ts):
    
    s = ts // 1000000000
    ms = ts % 1000000000
    m = s // 60
    s = s % 60
    h = m // 60
    m = m % 60
    
    return '{0:02d}:{1:02d}:{2:02d}.{3:09d}'.format(h, m, s, ms)

def determineFormat(timecodes):
    if rat.match(timecodes) != None or re.match('\d+',timecodes) != None:
        return 1
    elif v1re.match(linecache.getline(timecodes,1)):
        return 2
    elif v2re.match(linecache.getline(timecodes,1)):
        return 3
    else:
        return 0

def Ts(fn,tc):
    
    # CFR
    if tc[1] == 1:
        fps = rat.search(tc[0]).groups() if rat.search(tc[0]) else [re.search('(\d+)',tc[0]).group(0),'1']
        ts = round((1000 * fn * float(fps[1])) / float(fps[0]),6)
    # VFR
    elif tc[1] == 2 or tc[1] == 3:
        ts = linecache.getline(tc[0],fn+2)
    else:
        sys.exit("Couldn't get timestamps")
        
    
    return int(float(ts)*1000000)

def generateChap(start, end, chapter, type):
    matroskaXml = """
		<ChapterAtom>
			<ChapterTimeStart>{0}</ChapterTimeStart>
			<ChapterTimeEnd>{1}</ChapterTimeEnd>
			<ChapterDisplay>
				<ChapterString>Chapter {2:02d}</ChapterString>
				<ChapterLanguage>{3}</ChapterLanguage>
			</ChapterDisplay>
		</ChapterAtom>
"""[1:].format(start,end,chapter,"eng")
    ogmTxt = 'CHAPTER{0:02d}={1}\nCHAPTER{0:02d}NAME=Chapter {0:02d}\n'.format(chapter,start[:-6])
    if type == 'MKV':
        return matroskaXml
    elif type == 'OGM':
        return ogmTxt

if __name__ == '__main__':
    main()
